import argparse
import os
import sys
sys.path.append(os.getcwd())
sys.path.append("..")
import pandas as pd
from math import ceil
from tqdm import tqdm
from src.utils import mutant_filter

def top_x_factory(df_truth, df_predict, x):
    n = ceil(len(df_truth.index) * x)
    top_x_mutants = df_truth.sort_values("score").mutant[-n:]
    pred_x_mutants = df_predict.sort_values("score").mutant[-n:]
    real_top = set(pred_x_mutants) & set(top_x_mutants)
    return len(real_top) / len(top_x_mutants)

def process(args):
    proteins = [p.replace("\n", "") for p in open(args.trg_proteins, "r").readlines()]
    # aggrate the results
    info = {}
    for p in proteins:
        if args.sep:
            name, cat = p.split(args.sep)[0], p.split(args.sep)[1]
            if name not in info.keys():
                info[name] = [cat]
            else:
                info[name].append(cat)
        else:
            info[p] = [p]
    
    cols = ["name", "count"]
    cols.extend(args.model_locate)
    df = pd.DataFrame(columns=cols)
    
    
    for rate in args.x:
        bar = tqdm(args.model_locate)
        for model_name in bar:
            bar.set_postfix_str(model_name)
            names, scores, lens = [], [], []
            for p, i in info.items():
                for c in i:
                    try:
                        ground_truth = pd.read_table(f"{args.dataset_input}/{p}/experiments/{p}-{c}.tsv")
                        protein_res = pd.read_table(f"{args.dataset_input}/{p}/predictions/{p}-{c}.{model_name}.tsv")
                        
                    except:
                        try:
                            ground_truth = pd.read_table(f"{args.dataset_input}/{p}-{c}/experiments/{p}-{c}.tsv")
                            protein_res = pd.read_table(f"{args.dataset_input}/{p}-{c}/predictions/{p}-{c}.{model_name}.tsv")
                        except:
                            ground_truth = pd.read_table(f"{args.dataset_input}/{p}/experiments/{p}.tsv")
                            protein_res = pd.read_table(f"{args.dataset_input}/{p}/predictions/{p}.{model_name}.tsv")
                    protein_res = mutant_filter(protein_res, args.mutant_site).dropna()
                    ground_truth = mutant_filter(ground_truth, args.mutant_site).loc[protein_res.index]
                    assert len(ground_truth) == len(protein_res), f"{p}-{c} {len(ground_truth)} {len(protein_res)}"
                    score = top_x_factory(ground_truth, protein_res, rate)
                    if args.sep:
                        names.append(f"{p}{args.sep}{c}")
                    else:
                        names.append(f"{p}")
                    scores.append(score)
                    lens.append(len(ground_truth))
            
            df["count"] = lens
            df["name"] = names
            df[model_name] = scores
        out_path = os.path.join(args.output_dir, args.trg_proteins.split("/")[-1].split(".")[0]+f"_{args.mutant_site}_top{rate}.csv")
        print(f"Saving to {out_path}")
        df.to_csv(out_path, index=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--x", type=float, default=0.05, nargs="+")
    parser.add_argument("--trg_proteins", type=str, default="/nvme/tyang/workspace/hestia/private.txt")
    parser.add_argument("--dataset_input", type=str, default="/nvme/tyang/workspace/Proteus/data/common_proteins")
    parser.add_argument("--mutant_site", type=int, default=0)
    parser.add_argument("--sep", type=str)
    parser.add_argument("--model_locate", type=str, nargs="+", default=["esm1v_avg"])
    parser.add_argument("--output_dir", type=str, default="/nvme/tyang/workspace/hestia")
    args = parser.parse_args()
    
    process(args)